import numpy as np
import cv2
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
import pickle

def load_image(image_path, segmentation_path):
    """加载图像和分区图像"""
    try:
        image = cv2.imread(image_path, 0)
        segmentation = cv2.imread(segmentation_path, 0)
        if image is None or segmentation is None:
            raise FileNotFoundError("图像或分区图像未找到")
        return image, segmentation
    except Exception as e:
        print(f"加载图像时出错: {e}")
        raise

def preprocess_data(image, segmentation):
    """预处理图像和分区数据"""
    X = image.reshape(-1, 1)
    y = segmentation.reshape(-1)
    print("完成从砂岩截面图1及其对应分区中获取X和y")
    return X, y


def train_model(X_train, y_train):
    """训练随机森林分类器模型"""
    clf = RandomForestClassifier()
    clf.fit(X_train, y_train)
    print("完成train_test_split")
    print("完成随机森林模型clf的训练")
    return clf


def save_model(clf, model_path):
    """保存模型到硬盘"""
    with open(model_path, "wb") as file:
        pickle.dump(clf, file)


def main():
    image_path = r"D:/tuxiang/stone/Sandstone_1.tif"
    segmentation_path = r"D:/tuxiang/stone/Sandstone_1_segment.tif"
    model_path = r"D:/tuxiang/stone/clf.pkl"

    # 加载图像和分区图像
    image, segmentation = load_image(image_path, segmentation_path)
    print("图像形状:", image.shape)
    # 预处理数据
    X, y = preprocess_data(image, segmentation)

    # 分割数据集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # 训练模型
    clf = train_model(X_train, y_train)

    # 在测试集上进行预测并输出准确率
    accuracy = clf.score(X_test, y_test)
    print(f"模型准确率: {accuracy}")

    # 保存模型到硬盘
    save_model(clf, model_path)
    print("已保存随机森林模型clf到硬盘")

if __name__ == "__main__":
    main()